/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import java.util.UUID;
import java.util.concurrent.CountDownLatch;

import org.slf4j.Logger;
import org.sss.common.Bucket;
import org.sss.common.DataReceiver;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.Partitioner;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.datatype.Packable;

/**
 * Server side DataPutter.
 *
 * NOTE: Thid class for test.
 */
public class PutTask<K extends Packable, V extends Packable> extends Task {
  private static Logger logger = SssServer.getLogger(PutTask.class);
  private final TaskOutput outlet;
  private final CountDownLatch writeLatch;
  private Throwable exception = null;

  public static <K extends Packable, V extends Packable>
  PutTask<K, V> create(SssServerResource sssResource,                                   
                       Configuration conf,                                              
                       GroupID outputGID,                                               
                       Class<? extends Partitioner> partitionerClass,                   
                       Class<? extends Packable> outputKeyClass,                        
                       Class<? extends Packable> outputValueClass) throws SssException {
    return new PutTask<K, V>(sssResource, conf, outputGID, partitionerClass, outputKeyClass, outputValueClass);
  }

  private PutTask(SssServerResource sssResource,
                  Configuration conf,
                  GroupID outputGID,
                  Class<? extends Partitioner> partitionerClass,
                  Class<? extends Packable> outputKeyClass,
                  Class<? extends Packable> outputValueClass) throws SssException {
    super(new TaskSet(), UUID.randomUUID(), sssResource, conf, null, new Input(null, outputKeyClass, outputValueClass));
    this.outlet = getOutlets().create(outputGID,
        outputKeyClass, outputValueClass, null, partitionerClass,
        128 * 1024,
        (int)2,
        (int)2 * 1024 * 1024);
    this.writeLatch = new CountDownLatch(1);
  }

  public void write(K key, V value) throws SssException {
    DataReceiver<Bucket> cache = outlet.getTLPutter(1);
    try {
      cache.put(new Bucket(key, value));
    }
    finally {
      cache.close();
    }
  }

  public void close() throws SssException {
    outlet.close();
    try {
      writeLatch.await();
    }
    catch (InterruptedException e) {
      logger.error("caught an exception - {}", e);
    }
    if (exception != null) {
      throw new SssException(exception);
    }
  }

  @Override
  public synchronized void ioWriteComplete(Throwable e) {
    logger.debug("ioWriteComplete.");
    this.exception = e;
    writeLatch.countDown();
  }
}
